# states/task_state.py（完整版本，含 TaskStateManager 类）
from typing import Optional, Dict, Any
from datetime import datetime
import json
from loguru import logger

# 导入核心依赖
from utils.logger import Logger
from utils.exceptions import RedisError, LockAcquireError, ParamError
from utils.constants import (
    CACHE_EXPIRE_TASK_STATE,
    TASK_STATE_UPDATE_CHANNEL,
    TaskState,
    AgentType,
    ErrorCode
)
from middleware.redis_adapter import get_redis_adapter

class TaskStateManager:  # 确保该类存在且正确定义
    """任务状态管理器（核心类，管理任务全生命周期状态）"""
    def __init__(
        self,
        task_id: str,
        redis_client=None  # 兼容原有接口，现在使用RedisAdapter
    ):
        self.task_id = task_id
        self.redis_key = f"task:state:{self.task_id}"  # 任务状态Redis键
        self.lock_key = f"task:lock:{self.task_id}"    # 分布式锁键
        self.redis_client = redis_client or get_redis_adapter()
        self.logger = Logger.get_logger(task_id=self.task_id, agent_name="task_state_manager")

        # 初始化任务基础状态（未存在则创建）
        if not self.redis_client.hgetall(self.redis_key):
            self._init_task_state()

    def _init_task_state(self) -> bool:
        """初始化任务状态（默认待执行）"""
        init_data = {
            "task_id": self.task_id,
            "state": TaskState.PENDING.value,  # 引用 TaskState 枚举
            "created_at": str(int(datetime.now().timestamp())),
            "updated_at": str(int(datetime.now().timestamp())),
            "agent_type": "",  # 执行该任务的Agent类型
            "agent_instance_id": "",  # 执行该任务的Agent实例ID
            "input_params": json.dumps({}, ensure_ascii=False),  # 任务输入参数
            "output_result": json.dumps({}, ensure_ascii=False),  # 任务输出结果
            "error_code": "",  # 错误码（参考 ErrorCode 枚举）
            "error_msg": "",  # 错误信息
            "retry_count": "0",  # 重试次数
            "max_retries": "3"  # 最大重试次数
        }
        try:
            self.redis_client.hsetall(self.redis_key, init_data)
            self.redis_client.expire(self.redis_key, CACHE_EXPIRE_TASK_STATE)
            self.logger.debug(f"任务状态初始化成功 task_id={self.task_id}")
            return True
        except Exception as e:
            self.logger.error(f"任务状态初始化失败 task_id={self.task_id}: {str(e)}")
            raise RedisError(message="Task state init failed", context={"task_id": self.task_id, "error": str(e)}) from e

    def update_state(
        self,
        state: str,
        agent_type: Optional[str] = None,
        agent_instance_id: Optional[str] = None,
        input_params: Optional[Dict[str, Any]] = None,
        output_result: Optional[Dict[str, Any]] = None,
        error_code: Optional[str] = None,
        error_msg: Optional[str] = None
    ) -> bool:
        """更新任务状态（带分布式锁）"""
        # 校验状态合法性（用 TaskState 枚举约束）
        if state not in [s.value for s in TaskState]:
            raise ParamError(message="Invalid task state", context={"task_id": self.task_id, "state": state})

        if not self.redis_client.acquire_lock(self.lock_key, timeout=3):
            self.logger.error(f"获取任务锁失败 task_id={self.task_id}")
            raise LockAcquireError(message="Task lock acquire failed", context={"task_id": self.task_id})

        try:
            # 构造更新数据
            update_data = {
                "state": state,
                "updated_at": str(int(datetime.now().timestamp()))
            }
            if agent_type:
                # 校验Agent类型合法性
                if agent_type not in [t.value for t in AgentType]:
                    raise ParamError(message="Invalid agent type", context={"task_id": self.task_id, "agent_type": agent_type})
                update_data["agent_type"] = agent_type
            if agent_instance_id:
                update_data["agent_instance_id"] = agent_instance_id
            if input_params:
                update_data["input_params"] = json.dumps(input_params, ensure_ascii=False)
            if output_result:
                update_data["output_result"] = json.dumps(output_result, ensure_ascii=False)
            if error_code:
                update_data["error_code"] = error_code
            if error_msg:
                update_data["error_msg"] = error_msg

            # 更新重试次数（失败/重试状态时+1）
            if state in [TaskState.FAILED.value, TaskState.RETRYING.value]:
                current_retry = int(self.redis_client.hget(self.redis_key, "retry_count") or "0")
                update_data["retry_count"] = str(current_retry + 1)

            # 执行更新
            self.redis_client.hsetall(self.redis_key, update_data)
            self.redis_client.expire(self.redis_key, CACHE_EXPIRE_TASK_STATE)  # 刷新过期时间

            # 发布状态更新消息
            self._publish_state_update(update_data)
            self.logger.info(f"任务状态更新成功 task_id={self.task_id} | state={state}")
            return True
        except ParamError:
            raise
        except Exception as e:
            self.logger.error(f"任务状态更新失败 task_id={self.task_id}: {str(e)}")
            raise RedisError(message="Task state update failed", context={"task_id": self.task_id, "error": str(e)}) from e
        finally:
            self.redis_client.release_lock(self.lock_key)

    def get_state(self) -> Dict[str, Any]:
        """获取任务完整状态"""
        try:
            data = self.redis_client.hgetall(self.redis_key)
            if not data:
                self.logger.warning(f"任务状态不存在，自动初始化 task_id={self.task_id}")
                self._init_task_state()
                data = self.redis_client.hgetall(self.redis_key)

            # 解析JSON字段
            parsed_data = {
                "task_id": data.get("task_id", ""),
                "state": data.get("state", TaskState.PENDING.value),
                "created_at": int(data.get("created_at", 0)),
                "updated_at": int(data.get("updated_at", 0)),
                "agent_type": data.get("agent_type", ""),
                "agent_instance_id": data.get("agent_instance_id", ""),
                "input_params": json.loads(data.get("input_params", "{}")),
                "output_result": json.loads(data.get("output_result", "{}")),
                "error_code": data.get("error_code", ""),
                "error_msg": data.get("error_msg", ""),
                "retry_count": int(data.get("retry_count", 0)),
                "max_retries": int(data.get("max_retries", 3))
            }
            return parsed_data
        except Exception as e:
            self.logger.error(f"获取任务状态失败 task_id={self.task_id}: {str(e)}")
            raise RedisError(message="Task state get failed", context={"task_id": self.task_id, "error": str(e)}) from e

    def _publish_state_update(self, update_data: Dict[str, str]) -> None:
        """发布任务状态更新消息（引用 TASK_STATE_UPDATE_CHANNEL）"""
        try:
            message = {
                "task_id": self.task_id,
                "state": update_data["state"],
                "updated_at": update_data["updated_at"],
                "agent_type": update_data.get("agent_type", ""),
                "agent_instance_id": update_data.get("agent_instance_id", "")
            }
            self.redis_client.publish(TASK_STATE_UPDATE_CHANNEL, message)  # 确保引用正确
            self.logger.debug(f"发布任务状态消息 task_id={self.task_id} | state={update_data['state']}")
        except Exception as e:
            self.logger.warning(f"发布任务状态消息失败 task_id={self.task_id}: {str(e)}")

    def is_max_retry(self) -> bool:
        """判断是否已达到最大重试次数"""
        state_data = self.get_state()
        return state_data["retry_count"] >= state_data["max_retries"]

# 测试代码（验证类是否正常）
if __name__ == "__main__":
    import sys
    sys.path.append("..")

    # 初始化任务状态管理器
    task_manager = TaskStateManager(task_id="test_task_001")
    print("任务初始状态:", task_manager.get_state()["state"])

    # 更新任务状态为执行中
    task_manager.update_state(
        state=TaskState.RUNNING.value,
        agent_type=AgentType.TEXT_QA.value,
        input_params={"question": "测试问题"}
    )
    print("任务执行中状态:", task_manager.get_state()["state"])

    # 更新任务状态为成功
    task_manager.update_state(
        state=TaskState.SUCCESS.value,
        output_result={"answer": "测试答案"}
    )
    print("任务成功状态:", task_manager.get_state()["state"])